Streamed cohort diff behind a flag, hardened for 10-40M member cohorts (supersedes #38, includes #39) - #40
Merged
Conversation
A set difference over the full old/new member sets executes as a single blocking command on one single-threaded shard - multi-second for multi-million member cohorts - stalling every concurrent command on that shard for its whole duration (replicas included, since they re-execute the replicated command). At 10M members the diff reliably exceeds the default 10s command timeout, failing the refresh outright. Stream both sets via SSCAN in bounded chunks instead and diff client-side, hash-partitioned to cap proxy memory, so the largest single Redis command issued during a cohort update is one SSCAN page regardless of cohort size. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A cohort version's member set is written to Redis with no TTL, and only receives one when a *successor* refresh completes successfully. Any refresh that dies mid-flight (Redis command timeout, pod restart, CohortTooLargeException) therefore strands the full member set in Redis forever. For multi-million-member cohorts this leaks ~0.5 GB per failed refresh, all pinned to the same cluster slot by the cohort hash tag. In one production cluster we found 58 orphaned version sets (~25 GB) for a single 9.8M-member cohort, accumulated over two months of intermittent refresh failures. Fix, in three parts: - Pending version sets self-clean: apply a 24h TTL to the new version's member set as soon as ingestion starts writing it, and PERSIST it when the version is promoted (description published). An aborted refresh now leaves nothing behind once the TTL fires. - Temp diff keys get the same backstop TTL right after SDIFFSTORE, in case the process dies before the finally-block DELs run. - Retire the previous version only after successful promotion. Previously the EXPIRE of the old (still published) version set ran in a finally block, so a refresh that failed mid-update expired the *live* version, breaking cohort reads until the next successful refresh. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…orts
The streamed client-side diff is now opt-in via
AMPLITUDE_COHORT_STREAMED_DIFF_ENABLED (default false); the server-side
SDIFFSTORE diff remains the default path with a command stream identical
to before, plus the pending-version backstop TTLs from the previous
commit armed immediately after each SDIFFSTORE (not after both, so
addedKey is not left unprotected through the second, multi-second
blocking command). Two knobs are exposed (and validated > 0 at startup)
for tuning at 10M+ member scale, where the 2M default partition size
costs one full SSCAN pass over both member sets per partition:
AMPLITUDE_COHORT_DIFF_PARTITION_MAX_MEMBERS and
AMPLITUDE_COHORT_DIFF_SCAN_CHUNK_SIZE.
Hardening for the streamed path, whose diffs can run for minutes on
multi-million member cohorts:
- Renew the cohort-loading lock every 512 scan pages/flushes so a diff
outlasting the lock TTL (900s) doesn't let a second instance start a
concurrent load of the same cohort; abort the refresh (retried next
sync cycle) if renewal fails rather than continue unlocked. New
Redis.renewLock() uses an atomic compare-and-expire mirroring
releaseLock().
- Cross-check each scan's raw returned count against SCARD and abort
the refresh on a shortfall: SSCAN silently skips members when its
cursor is invalidated mid-scan (cluster failover/slot migration), and
a truncated scan here becomes false removals of unchanged users that
no later diff repairs. A distinct-count check across partitions
additionally catches duplicate-masked skips on the existing set.
- If the existing members key is missing at diff time, apply all new
members as additions with no removals -- SDIFFSTORE semantics for
that state -- under the same scan guard, and log an error instead of
silently publishing.
Also fixed in passing: lock values in RedisConnection and
RedisClusterConnection were built with ${'$'} escapes, producing the
same literal string on every instance and making releaseLock's
compare-and-delete vacuous across pods.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The kdoc predated the flag gating and still described the streamed diff as having replaced SDIFFSTORE; with the flag off (the default) the SDIFFSTORE path is what runs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Cursor bugbot on PR #39 flagged that persist(newCohortKey) ran before the description hset: a refresh dying between the two left a persisted (no-TTL) member set with the old description still published -- the pending TTL's self-clean guarantee was defeated exactly one command before it mattered. (In practice the next retry re-arms the TTL on the same key, so the leak only stuck when the upstream version was superseded first -- but the window was real.) Simply swapping the order would trade that leak for something worse: a crash after hset but before persist would leave the LIVE version with a 24h fuse and no code path to clear it (the next sync short-circuits on lastModified), breaking reads a day later. So this does both halves: - complete() now publishes first and persists immediately after, so a refresh that dies anywhere before publish always leaves a set that self-cleans via its pending TTL. - CohortLoader clears any TTL found on the published version's member set on every not-modified sync cycle (new O(1) CohortStorage.ensureCurrentVersionPersisted), repairing the one-command publish->persist crash window within one sync interval, instead of never. This also defuses the stray-fuse race where a lock-breaching concurrent writer re-arms the TTL on an already-promoted set. New tests: a fault-injected hset failure proves the pending TTL survives a crash at publish; ensureCurrentVersionPersisted clears a stray TTL on the published version. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Jul 21, 2026
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Partition count ignores Redis cardinality
- Changed partition count calculation to use existingCard and newCard from SCARD instead of description sizes existingSize and newSize, ensuring memory bounds are enforced based on actual Redis data.
Or push these changes by commenting:
@cursor push 9c51505bdf
Preview (9c51505bdf)
diff --git a/core/src/main/kotlin/cohort/CohortStorage.kt b/core/src/main/kotlin/cohort/CohortStorage.kt
--- a/core/src/main/kotlin/cohort/CohortStorage.kt
+++ b/core/src/main/kotlin/cohort/CohortStorage.kt
@@ -362,7 +362,7 @@ internal class RedisCohortStorage(
applyAllAdditions(newCohortKey, description, existingSize, newCard, cohortIdSet)
return
}
- val partitions = ((maxOf(existingSize, newSize, 1) - 1) / diffPartitionMaxMembers) + 1
+ val partitions = ((maxOf(existingCard, newCard, 1L) - 1) / diffPartitionMaxMembers).toInt() + 1
var addedCount = 0L
var removedCount = 0L
var existingDistinct = 0L
@@ -362,7 +362,7 @@ internal class RedisCohortStorage(
applyAllAdditions(newCohortKey, description, existingSize, newCard, cohortIdSet)
return
}
- val partitions = ((maxOf(existingSize, newSize, 1) - 1) / diffPartitionMaxMembers) + 1
+ val partitions = ((maxOf(existingCard, newCard, 1L) - 1) / diffPartitionMaxMembers).toInt() + 1
var addedCount = 0L
var removedCount = 0L
var existingDistinct = 0LYou can send follow-ups to the cloud agent here.
Reviewed by Cursor Bugbot for commit f09c9f9. Configure here.
Cursor bugbot on #40: partition count used the description-reported sizes (prev.size / finalSize) while the scan guards two lines above already fetch the authoritative SCARD cardinalities. A crashed ingest can leave a version key holding more members than its published size; partitioning by the smaller number lets a single in-memory partition exceed diffPartitionMaxMembers -- the bound partitioning exists to enforce. Partition count now comes from max(existingCard, newCard). The newSize parameter became unused and is removed. New test publishes a version whose description size understates the key's cardinality and verifies diff behavior (removals applied, retained/added memberships correct, residue members untouched -- matching SDIFFSTORE semantics). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Content-neutral merge: this branch already contained #39's changes (cherry-picked) and the backstop-ordering fix, so every conflict resolved to the branch side; the tree is identical to the pre-merge commit. Verified: full test suite and ktlint green, zero content diff vs 7632c55. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
kyeh-amp
self-requested a review
July 21, 2026 23:38
kyeh-amp
approved these changes
Jul 21, 2026
vaibhav-jain-exp
added a commit
that referenced
this pull request
Jul 22, 2026
* Add cohort_streamed_diff_enabled input to perf test workflow Adds a workflow_dispatch input that is forwarded to the sdk-tests repository_dispatch payload so a manual perf test run can enable AMPLITUDE_COHORT_STREAMED_DIFF_ENABLED (default false, added in #40) on the perf-test deployment. Push-triggered runs always send false. * Pass workflow inputs to run scripts via env to fix semgrep shell-injection findings
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Builds on @adrianliu0815's #38 and #39 (both authored commits preserved on this branch) and adds the gating and hardening we need before the streamed diff can be on a release path. Supersedes #38; if #39 merges first, the identical hunks here resolve cleanly.
From #38 (unchanged): the client-side streamed cohort diff — SSCAN both member-set versions in hash partitions and compute adds/removals in the proxy, so no O(cohort-size) blocking command ever runs on a Redis shard.
From #39 (cherry-picked): pending-version 24h TTL at ingestion, PERSIST at promotion, and retiring the previous version only after successful promotion — so failed refreshes can neither strand full member sets in Redis nor expire the still-published live version.
Added here:
AMPLITUDE_COHORT_STREAMED_DIFF_ENABLED(defaultfalse). The SDIFFSTORE path remains the default with a command stream identical to main, plus Prevent orphaned cohort version sets when a refresh fails mid-flight #39's temp-key backstop TTLs armed immediately after each SDIFFSTORE (previouslyaddedKeysat unprotected through the second, multi-second blocking command — the likeliest crash point).AMPLITUDE_COHORT_DIFF_PARTITION_MAX_MEMBERS(default 2M) andAMPLITUDE_COHORT_DIFF_SCAN_CHUNK_SIZE(default 1000). At 40M members the defaults cost ~1.6M sequential SSCAN round trips per refresh (~10-45 min); partition 8-10M + chunk 5000 brings that to ~1-4 min. NotebooleanEnvonly acceptstrue—1/yessilently mean off.Redis.renewLock()(compare-and-expire, mirroringreleaseLock); the diff renews the 900s cohort-loading lock every 512 scan pages and aborts fail-closed if renewal fails, so a diff outlasting the lock TTL can't run concurrently with a second instance's load of the same cohort.Redis.scard(); every SSCAN pass cross-checks raw returned count against cardinality and aborts on shortfall. A cluster failover or slot migration can silently invalidate an SSCAN cursor mid-scan; on this code path a truncated scan becomes false removals of unchanged users that no later diff repairs. A cross-partition distinct-count check additionally catches duplicate-masked skips on the existing set.complete()now publishes the description before PERSIST, so a refresh dying pre-publish always self-cleans via the pending TTL; and the inverse crash window (live set left with a 24h fuse that nothing would ever clear, which a plain order-swap would create) is repaired within one sync cycle by an O(1)ensureCurrentVersionPersistedon the not-modified loader path.${'$'}escapes, producing the same literal string on every instance and makingreleaseLock's compare-and-delete vacuous across pods.Testing
CohortStorageTest14/14 (union of Replace SDIFFSTORE cohort diff with streamed client-side diff #38/Prevent orphaned cohort version sets when a refresh fails mid-flight #39 tests plus: flag-off SDIFFSTORE update path, >1000 flush-threshold streamed update, missing-key fallback, knob validation, fault-injected crash-at-publish TTL survival, stray-TTL self-heal); full./gradlew test ktlintCheckgreen.expirearmed acohortSyncIntervalMillisfuse on it).Rollout guidance for large-cohort deployments
Flag on with partition 8-10M and chunk 5000 for 10-40M member cohorts; keep the cohort sync interval at or above expected diff duration; watch per-cohort diff duration and added/removed counts (a full-size "added" spike with zero removals is the signature of a missing-previous-key fallback).
Note
High Risk
Changes cohort refresh semantics, per-user membership writes, and distributed locking in Redis; incorrect diff or lock handling could corrupt memberships or allow concurrent loads on multi-million-member cohorts.
Overview
Adds an opt-in streamed client-side cohort diff (
AMPLITUDE_COHORT_STREAMED_DIFF_ENABLED, default off) so very large refreshes avoid blockingSDIFFSTOREon a Redis shard; when disabled, behavior stays on the existing server-side diff path with safer temp-key TTL handling.Streamed diff (when enabled): partitions member sets by hash, diffs via chunked
SSCAN+SCARDcompleteness checks, pipelines membership updates, renews the cohort loading lock during long scans, and degrades to “all additions” if the previous members key is missing.Promotion / recovery:
complete()now publishes the description beforePERSISTon the new member set;ensureCurrentVersionPersistedon unchanged sync cycles clears a stray TTL on the live set. Redis:scard,renewLock, per-instance lock values (fixes broken${'$'}lock strings), and env-tunable partition/chunk sizes validated at startup.Tests cover both diff paths, crash-at-publish TTL, and TTL self-heal.
Reviewed by Cursor Bugbot for commit 1556659. Bugbot is set up for automated code reviews on this repo. Configure here.